# -*- coding: utf-8 -*-

import sys
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QTextEdit, QGridLayout, QLabel
from PyQt5.QtGui import QFont, QPixmap
from PyQt5.QtCore import QEvent, Qt
from flow.label.labeled_backend_processor import LabeledBackendProcessor
from pathlib import Path
import pickle

QApplication.addLibraryPath("./PyQt5/Qt/plugins")


class LabelWindow(QWidget):
    def __init__(self, src_folder, dst_folder):
        super(LabelWindow, self).__init__()
        self.status_line_edit = QTextEdit()
        self.previous_image_button = QPushButton("返回上一张")
        self.next_image_button = QPushButton("进入下一张")
        self.image_label = QLabel()
        self.backend_processor = LabeledBackendProcessor(src_folder, dst_folder)

        self.previous_image_button.clicked.connect(self.previous_image_button_pressed)
        self.next_image_button.clicked.connect(self.next_image_button_pressed)

        self.init_ui()
        self.set_up_child_widgets_property()
        self.setGeometry(300, 300, 1000, 550)
        self.begin_label()
        self.show()

    def begin_label(self):
        processed_count, left_count = self.backend_processor.load_an_image()
        if processed_count > 0:
            self.previous_image_button.setDisabled(False)
        if left_count > 0:
            self.next_image_button.setDisabled(False)
            self.image_label.setPixmap(QPixmap(self.backend_processor.draw_image()))
        self.show_label_status(processed_count, left_count)

    def init_ui(self):
        font = QFont("宋体", 15)
        font.setBold(True)
        self.status_line_edit.setFixedHeight(40)
        self.previous_image_button.setFixedHeight(40)
        self.previous_image_button.setFont(font)
        self.next_image_button.setFixedHeight(40)
        self.next_image_button.setFont(font)
        grid_layout = QGridLayout()
        grid_layout.addWidget(self.status_line_edit, 0, 0, 1, 8)
        grid_layout.addWidget(self.previous_image_button, 0, 8, 1, 2)
        grid_layout.addWidget(self.next_image_button, 0, 10, 1, 2)
        grid_layout.setSpacing(15)
        grid_layout.addWidget(self.image_label, 2, 0, 2, 12)
        self.setLayout(grid_layout)
        self.setWindowTitle("客流标注工具")

    def set_up_child_widgets_property(self):
        self.status_line_edit.setDisabled(True)
        self.previous_image_button.setDisabled(True)
        self.next_image_button.setDisabled(True)
        self.image_label.setScaledContents(True)
        self.image_label.installEventFilter(self)
        self.image_label.setMouseTracking(True)
        self.image_label.setFocus()

    def show_label_status(self, processed_count, left_count):
        text = "<span style=\"font-size:18pt; font-weight:600; color:#ff0000;\">"
        text += "正在标注。。。"
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += "已经标注了 "
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#0000ff;\">"
        text += "%d" % processed_count
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += " 张图片"
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += ", 剩余 "
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#0000ff;\">"
        text += "%d" % left_count
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += " 张图片"
        text += "</span>"
        self.status_line_edit.setText(text)

    def get_relative_position(self, event):
        x = event.x()
        y = event.y()
        relative_position = [x / self.image_label.width(), y / self.image_label.height()]
        return relative_position

    def eventFilter(self, source, event):
        if source is not self.image_label:
            return QWidget.eventFilter(self, source, event)
        if event.type() == QEvent.MouseMove:
            relative_position = self.get_relative_position(event)
            self.image_label.setPixmap(
                QPixmap(self.backend_processor.draw_image(relative_position,
                                                          draw_dashed_line=True,
                                                          draw_mouse_move_rect=True)))
        if event.type() == QEvent.MouseButtonPress and event.button() == Qt.LeftButton:
            relative_position = self.get_relative_position(event)
            self.backend_processor.set_left_top_point(relative_position)
        if event.type() == QEvent.MouseButtonRelease and event.button() == Qt.LeftButton:
            relative_position = self.get_relative_position(event)
            self.image_label.setPixmap(
                QPixmap(self.backend_processor.draw_image(relative_position,
                                                          draw_dashed_line=True,
                                                          draw_mouse_release_rect=True)))
        if event.type() == QEvent.MouseButtonRelease and event.button() == Qt.RightButton:
            self.image_label.setPixmap(
                QPixmap(self.backend_processor.right_mouse_button_event(event.pos().x() / self.image_label.width(),
                                                                        event.pos().y() / self.image_label.height())))
        if event.type() == QEvent.KeyPress:
            if event.key() == Qt.Key_Escape:
                self.next_image_button_pressed()
            if event.key() == Qt.Key_D:
                self.delete_image()
        return QWidget.eventFilter(self, source, event)

    def delete_image(self):
        processed_count, left_count = self.backend_processor.delete_current_image()
        if left_count == 0:  # 所有图片已经标注完毕
            self.next_image_button.setDisabled(True)
            self.image_label.clear()
        else:  # 显示下一张图片
            self.next_image_button.setDisabled(False)
            self.image_label.setPixmap(QPixmap(self.backend_processor.draw_image()))
        if processed_count > 0:
            self.previous_image_button.setDisabled(False)
        self.show_label_status(processed_count, left_count)
        self.image_label.setFocus()

    def previous_image_button_pressed(self):
        processed_count, left_count = self.backend_processor.previous_image_button_event()
        if processed_count == 0:  # 目标文件夹已经没有标注过的图片
            self.previous_image_button.setDisabled(True)
        if left_count > 0:
            self.next_image_button.setDisabled(False)
            self.image_label.setPixmap(QPixmap(self.backend_processor.draw_image()))
        self.show_label_status(processed_count, left_count)
        self.image_label.setFocus()

    def next_image_button_pressed(self):
        processed_count, left_count = self.backend_processor.next_image_button_event()
        if left_count == 0:  # 所有图片已经标注完毕
            self.next_image_button.setDisabled(True)
            self.image_label.clear()
        else:  # 显示下一张图片
            self.next_image_button.setDisabled(False)
            self.image_label.setPixmap(QPixmap(self.backend_processor.draw_image()))
        if processed_count > 0:
            self.previous_image_button.setDisabled(False)
        self.show_label_status(processed_count, left_count)
        self.image_label.setFocus()


def main():
    root_dir = Path('data')
    config_path = Path("config.data")
    if config_path.exists():
        with open(config_path, 'rb') as f:
            config = pickle.load(f)
        src_folder = root_dir / Path(config['head_rectangle_input'])
        dst_folder = root_dir / Path(config['head_rectangle_output'])
    else:
        src_folder = root_dir / "images"
        dst_folder = root_dir / "labeled"
    src_folder.mkdir(exist_ok=True)
    dst_folder.mkdir(exist_ok=True)
    app = QApplication(sys.argv)
    check_window = LabelWindow(str(src_folder), str(dst_folder))
    check_window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
